#!/usr/bin/env python
# coding: utf-8

import pdfplumber
from PyPDF2 import PdfReader

from utils.utils import get_path, logger


class DataProcess(object):

    def __init__(self, pdf_path):
        self.pdf_path = pdf_path
        self.data = []

    # 滑动窗口功能实现，其中fast代表当前遍历句子的index，slow代表每次窗口开始滑动的起点。默认窗口直接滑动的overlap是1个句子。
    def sliding_window(self, sentences, kernel=512, stride=1):
        sz = len(sentences)
        cur = ""
        fast = 0
        slow = 0
        while fast < len(sentences):
            sentence = sentences[fast]
            if len(cur + sentence) > kernel and (cur + sentence) not in self.data:
                self.data.append(cur + sentence + "。")
                cur = cur[len(sentences[slow] + "。") :]
                slow = slow + 1
            cur = cur + sentence + "。"
            fast = fast + 1

    #  数据过滤，根据当前的文档内容的item划分句子，然后根据max_seq划分文档块。
    def data_filter(self, line, header, pageid, max_seq=1024):

        sz = len(line)
        if sz < 6:
            return

        if sz > max_seq:

            if "■" in line:
                sentences = line.split("■")
            elif "•" in line:
                sentences = line.split("•")
            elif "\t" in line:
                sentences = line.split("\t")
            else:
                sentences = line.split("。")

            for subsentence in sentences:
                subsentence = subsentence.replace("\n", "")

                if len(subsentence) < max_seq and len(subsentence) > 5:
                    subsentence = (
                        subsentence.replace(",", "").replace("\n", "").replace("\t", "")
                    )
                    if subsentence not in self.data:
                        self.data.append(subsentence)
        else:
            line = line.replace("\n", "").replace(",", "").replace("\t", "")
            if line not in self.data:
                self.data.append(line)

    # 提取页头即一级标题
    def get_header(self, page):
        try:
            lines = page.extract_words()[::]
        except:
            return None
        if len(lines) > 0:
            for line in lines:
                if "目录" in line["text"] or ".........." in line["text"]:
                    return None
                if line["top"] < 20 and line["top"] > 17:
                    return line["text"]
            return lines[0]["text"]
        return None

    # 按照每页中块提取内容,并和一级标题进行组合,配合Document 可进行意图识别
    def parse_blocks(self, max_seq=1024):

        with pdfplumber.open(self.pdf_path) as pdf:

            for i, p in enumerate(pdf.pages):
                header = self.get_header(p)

                if header == None:
                    continue

                texts = p.extract_words(use_text_flow=True, extra_attrs=["size"])[::]

                squence = ""
                lastsize = 0

                for idx, line in enumerate(texts):
                    if idx < 1:
                        continue
                    if idx == 1:
                        if line["text"].isdigit():
                            continue
                    cursize = line["size"]
                    text = line["text"]
                    if text == "□" or text == "•":
                        continue
                    elif text == "警告！" or text == "注意！" or text == "说明！":
                        if len(squence) > 0:
                            self.data_filter(squence, header, i, max_seq=max_seq)
                        squence = ""
                    elif format(lastsize, ".5f") == format(cursize, ".5f"):
                        if len(squence) > 0:
                            squence = squence + text
                        else:
                            squence = text
                    else:
                        lastsize = cursize
                        if len(squence) < 15 and len(squence) > 0:
                            squence = squence + text
                        else:
                            if len(squence) > 0:
                                self.data_filter(squence, header, i, max_seq=max_seq)
                            squence = text
                if len(squence) > 0:
                    self.data_filter(squence, header, i, max_seq=max_seq)

    # 按句号划分文档，然后利用最大长度划分文档块
    def parse_one_page_with_rule(self, max_seq=512, min_len=6):
        for idx, page in enumerate(PdfReader(self.pdf_path).pages):
            page_content = ""
            text = page.extract_text()
            words = text.split("\n")
            for idx, word in enumerate(words):
                text = word.strip().strip("\n")
                if "...................." in text or "目录" in text:
                    continue
                if len(text) < 1:
                    continue
                if text.isdigit():
                    continue
                page_content = page_content + text
            if len(page_content) < min_len:
                continue
            if len(page_content) < max_seq:
                if page_content not in self.data:
                    self.data.append(page_content)
            else:
                sentences = page_content.split("。")
                cur = ""
                for idx, sentence in enumerate(sentences):
                    if (
                        len(cur + sentence) > max_seq
                        and (cur + sentence) not in self.data
                    ):
                        self.data.append(cur + sentence)
                        cur = sentence
                    else:
                        cur = cur + sentence

    #  滑窗法提取段落
    #  1. 把pdf看做一个整体,作为一个字符串
    #  2. 利用句号当做分隔符,切分成一个数组
    #  3. 利用滑窗法对数组进行滑动, 此处的
    def parse_all_page(self, max_seq=512, min_len=6):
        all_content = ""
        for idx, page in enumerate(PdfReader(self.pdf_path).pages):
            page_content = ""
            text = page.extract_text()
            words = text.split("\n")
            for idx, word in enumerate(words):
                text = word.strip().strip("\n")
                if "...................." in text or "目录" in text:
                    continue
                if len(text) < 1:
                    continue
                if text.isdigit():
                    continue
                page_content = page_content + text
            if len(page_content) < min_len:
                continue
            all_content = all_content + page_content
        sentences = all_content.split("。")
        self.sliding_window(sentences, kernel=max_seq)


if __name__ == "__main__":
    pdf_path = get_path("data/train_a.pdf")
    data_process = DataProcess(pdf_path)
    data_process.parse_blocks(max_seq=1024)
    data_process.parse_blocks(max_seq=512)
    logger.info(len(data_process.data))
    data_process.parse_all_page(max_seq=256)
    data_process.parse_all_page(max_seq=512)
    logger.info(len(data_process.data))
    data_process.parse_one_page_with_rule(max_seq=256)
    data_process.parse_one_page_with_rule(max_seq=512)
    logger.info(len(data_process.data))
    data = data_process.data
    out = open("all_text.txt", "w")
    for line in data:
        line = line.strip("\n")
        out.write(line)
        out.write("\n")
    out.close()
